import json
from typing import List, Tuple, Dict
from collections import defaultdict

# 确保导入所需的库
import pandas as pd
import re

# 导入开关解析函数
from restore_metirc import parse_action_switch

# 定义一个函数来标准化动作列表的顺序
def normalize_action_list(actions: List[Tuple[int, int, int]]):
    # 排序每个动作的首尾节点
    normalized_actions = [
        (min(from_bus, to_bus), max(from_bus, to_bus), status)
        for from_bus, to_bus, status in actions
    ]
    # 对整个列表排序
    normalized_actions.sort()
    return tuple(normalized_actions)

# 统计每种模式出现的数量
def count_action_patterns(jsonl_file_path: str) -> Dict[Tuple[Tuple[int, int, int], ...], int]:
    pattern_count = defaultdict(int)

    with open(jsonl_file_path, 'r', encoding='utf-8') as file:
        for line in file:
            # 去除行尾的换行符
            line = line.strip()

            # 跳过空行
            if not line:
                continue

            # 解析每一行的JSON数据
            data = json.loads(line)
            messages = data.get("messages", [])
            for message in messages:
                if message.get("role") == "assistant":
                    content = message.get("content", "")
                    actions, num_actions, find_action = parse_action_switch(content)
                    if find_action:
                        normalized_actions = normalize_action_list(actions)
                        pattern_count[normalized_actions] += 1

    return pattern_count

# 示例使用
if __name__ == "__main__":
    jsonl_file_path = "samples\\402算例_时序负荷数据\\train.jsonl"
    pattern_count = count_action_patterns(jsonl_file_path)

    # 打印出现的次数最多的前10种模式及其出现的次数
    sorted_patterns = sorted(pattern_count.items(), key=lambda x: x[1], reverse=True)[
        :10
    ]
    print("| 模式 | 出现次数 |")
    print("|------|----------|")
    for pattern, count in sorted_patterns:
        print(f"| {pattern} | {count} |")

    # 将模式及其出现的次数结果保存到CSV文件
    df = pd.DataFrame([
        {'Pattern': str(pattern), 'Count': count}
        for pattern, count in pattern_count.items()
    ])
    df.to_csv('action_patterns_count.csv', index=False, encoding='utf-8')
    print("结果已保存到 action_patterns_count.csv")
